#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脑电控制模型训练脚本
"""

import tensorflow as tf
import numpy as np
import yaml
import os
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
from eeg_dog_control_model import create_model, compile_model

def generate_synthetic_data(num_samples=1000, sequence_length=5, time_points=500):
    """
    生成合成数据用于演示
    在实际应用中，这里应该加载真实的脑电数据
    """
    print("📊 生成合成训练数据...")
    
    # 为每个控制类别生成数据
    data = []
    labels = []
    
    for class_id in range(4):  # 4个控制类别
        print(f"生成类别 {class_id} 的数据...")
        
        # 为每个类别生成特定模式的数据
        for i in range(num_samples // 4):
            # 生成基础脑电信号
            base_signal = np.random.randn(sequence_length, time_points)
            
            # 为不同类别添加特定模式
            if class_id == 0:  # 向前 - 添加高频成分
                pattern = np.sin(np.linspace(0, 10*np.pi, time_points)) * 0.3
                base_signal += pattern
            elif class_id == 1:  # 向后 - 添加低频成分
                pattern = np.sin(np.linspace(0, 5*np.pi, time_points)) * 0.3
                base_signal += pattern
            elif class_id == 2:  # 趴下 - 添加中频成分
                pattern = np.sin(np.linspace(0, 7*np.pi, time_points)) * 0.3
                base_signal += pattern
            else:  # 转圈 - 添加复杂模式
                pattern = np.sin(np.linspace(0, 8*np.pi, time_points)) * 0.3
                base_signal += pattern
            
            # 添加噪声
            noise = np.random.randn(sequence_length, time_points) * 0.1
            base_signal += noise
            
            # 重塑为模型输入格式
            sample = base_signal.reshape(sequence_length, time_points, 1, 1)
            data.append(sample)
            
            # 创建one-hot标签
            label = np.zeros(4)
            label[class_id] = 1
            labels.append(label)
    
    return np.array(data), np.array(labels)

def load_real_data(data_path):
    """
    加载真实脑电数据
    """
    print(f"📂 加载真实数据: {data_path}")
    
    try:
        # 这里应该实现真实数据的加载逻辑
        # 例如从npz文件或mat文件加载
        data = np.load(data_path)
        return data['X'], data['y']
    except Exception as e:
        print(f"❌ 加载真实数据失败: {e}")
        print("使用合成数据...")
        return None, None

def preprocess_data(X, y):
    """
    数据预处理
    """
    print("🔧 数据预处理...")
    
    # 数据标准化
    X_mean = np.mean(X, axis=(1, 2), keepdims=True)
    X_std = np.std(X, axis=(1, 2), keepdims=True)
    X_normalized = (X - X_mean) / (X_std + 1e-8)
    
    # 分割训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X_normalized, y, test_size=0.2, random_state=42, stratify=np.argmax(y, axis=1)
    )
    
    # 分割训练集和验证集
    X_train, X_val, y_train, y_val = train_test_split(
        X_train, y_train, test_size=0.2, random_state=42, stratify=np.argmax(y_train, axis=1)
    )
    
    print(f"训练集: {X_train.shape}")
    print(f"验证集: {X_val.shape}")
    print(f"测试集: {X_test.shape}")
    
    return X_train, X_val, X_test, y_train, y_val, y_test

def create_callbacks():
    """
    创建训练回调函数
    """
    callbacks = [
        # 早停
        tf.keras.callbacks.EarlyStopping(
            monitor='val_loss',
            patience=10,
            restore_best_weights=True,
            verbose=1
        ),
        
        # 学习率调度
        tf.keras.callbacks.ReduceLROnPlateau(
            monitor='val_loss',
            factor=0.5,
            patience=5,
            min_lr=1e-6,
            verbose=1
        ),
        
        # 模型检查点
        tf.keras.callbacks.ModelCheckpoint(
            'models/best_eeg_control_model.h5',
            monitor='val_accuracy',
            save_best_only=True,
            verbose=1
        ),
        
        # TensorBoard日志
        tf.keras.callbacks.TensorBoard(
            log_dir='logs/eeg_control',
            histogram_freq=1,
            write_graph=True
        )
    ]
    
    return callbacks

def train_model(model, X_train, y_train, X_val, y_val, config):
    """
    训练模型
    """
    print("🎯 开始训练模型...")
    
    # 创建回调函数
    callbacks = create_callbacks()
    
    # 训练模型
    history = model.fit(
        X_train, y_train,
        epochs=config['training']['epochs'],
        batch_size=config['training']['batch_size'],
        validation_data=(X_val, y_val),
        callbacks=callbacks,
        verbose=1
    )
    
    return history

def evaluate_model(model, X_test, y_test):
    """
    评估模型
    """
    print("📊 评估模型性能...")
    
    # 预测
    y_pred = model.predict(X_test)
    y_pred_classes = np.argmax(y_pred, axis=1)
    y_true_classes = np.argmax(y_test, axis=1)
    
    # 计算准确率
    accuracy = np.mean(y_pred_classes == y_true_classes)
    print(f"测试准确率: {accuracy:.4f}")
    
    # 分类报告
    class_names = ['向前', '向后', '趴下', '转圈']
    print("\n分类报告:")
    print(classification_report(y_true_classes, y_pred_classes, target_names=class_names))
    
    # 混淆矩阵
    cm = confusion_matrix(y_true_classes, y_pred_classes)
    print("\n混淆矩阵:")
    print(cm)
    
    # 绘制混淆矩阵
    plt.figure(figsize=(8, 6))
    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
    plt.title('混淆矩阵')
    plt.colorbar()
    plt.xticks(range(4), class_names)
    plt.yticks(range(4), class_names)
    
    # 添加数值标签
    thresh = cm.max() / 2
    for i in range(4):
        for j in range(4):
            plt.text(j, i, format(cm[i, j], 'd'),
                    ha="center", va="center",
                    color="white" if cm[i, j] > thresh else "black")
    
    plt.tight_layout()
    plt.savefig('confusion_matrix.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    return accuracy, y_pred

def plot_training_history(history):
    """
    绘制训练历史
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
    
    # 损失曲线
    ax1.plot(history.history['loss'], label='训练损失')
    ax1.plot(history.history['val_loss'], label='验证损失')
    ax1.set_title('模型损失')
    ax1.set_xlabel('轮次')
    ax1.set_ylabel('损失')
    ax1.legend()
    ax1.grid(True)
    
    # 准确率曲线
    ax2.plot(history.history['accuracy'], label='训练准确率')
    ax2.plot(history.history['val_accuracy'], label='验证准确率')
    ax2.set_title('模型准确率')
    ax2.set_xlabel('轮次')
    ax2.set_ylabel('准确率')
    ax2.legend()
    ax2.grid(True)
    
    plt.tight_layout()
    plt.savefig('training_history.png', dpi=300, bbox_inches='tight')
    plt.show()

def save_model(model, config):
    """
    保存模型
    """
    print("💾 保存模型...")
    
    # 创建模型目录
    os.makedirs('models', exist_ok=True)
    
    # 保存完整模型
    model.save('models/eeg_control_model.h5')
    
    # 保存为TFLite格式（用于部署）
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    tflite_model = converter.convert()
    
    with open('models/eeg_control_model.tflite', 'wb') as f:
        f.write(tflite_model)
    
    # 保存配置
    with open('models/model_config.yaml', 'w') as f:
        yaml.dump(config, f)
    
    print("✅ 模型保存完成！")

def main():
    """
    主函数
    """
    print("🧠 脑电控制模型训练")
    print("=" * 50)
    
    # 配置参数
    config = {
        'model': {
            'type': 'lightweight',
            'sequence_length': 5,
            'time_points': 500,
            'num_classes': 4
        },
        'training': {
            'epochs': 50,
            'batch_size': 16,
            'learning_rate': 0.001,
            'optimizer': 'adam'
        },
        'data': {
            'use_synthetic': True,  # 使用合成数据
            'num_samples': 1000,
            'train_split': 0.7,
            'val_split': 0.2,
            'test_split': 0.1
        }
    }
    
    # 创建模型
    print("🏗️ 创建模型...")
    model = create_model(config)
    model = compile_model(model, config)
    
    # 打印模型结构
    print("📊 模型结构:")
    model.summary()
    
    # 加载数据
    if config['data']['use_synthetic']:
        print("📈 使用合成数据...")
        X, y = generate_synthetic_data(
            num_samples=config['data']['num_samples'],
            sequence_length=config['model']['sequence_length'],
            time_points=config['model']['time_points']
        )
    else:
        print("📈 加载真实数据...")
        X, y = load_real_data('data/eeg_control_data.npz')
        if X is None:
            print("使用合成数据...")
            X, y = generate_synthetic_data(
                num_samples=config['data']['num_samples'],
                sequence_length=config['model']['sequence_length'],
                time_points=config['model']['time_points']
            )
    
    # 数据预处理
    X_train, X_val, X_test, y_train, y_val, y_test = preprocess_data(X, y)
    
    # 训练模型
    history = train_model(model, X_train, y_train, X_val, y_val, config)
    
    # 评估模型
    accuracy, predictions = evaluate_model(model, X_test, y_test)
    
    # 绘制训练历史
    plot_training_history(history)
    
    # 保存模型
    save_model(model, config)
    
    print("🎉 训练完成！")
    print(f"最终测试准确率: {accuracy:.4f}")

if __name__ == "__main__":
    main() 